
# 使用下面的 命令 运行此文件 的 main
# python -m transformer_code.learn_conf.batch_data

import torch

from ..model.part import subsequentMask


# ***************************************************************************************
# Usage : 创建一个 train 时的 batch，这个 batch 类存放了 src (encoder 输入), tgt (decoder 输入),
#         并创建相应的 mask，包括 padding mask, subsequent mask
#
# Note: 1. 创建过程中会 去掉 tgt 前后的 开始 和 结束 标志
#
# @input_param (src) : 作为 encoder 的输入，size : (batch_size, src_len),
#                      其中 batch_size 代表了这一个 batch 中有多少个 sequence, src_len 表示
#                      每个 seq 的长度，这个长度是经过 padding 的 (因为每个句子的长度都可能不一样)
# @input_param (tgt) : 作为 decoder 的输入，size : (batch_size, tgt_len) 同上
#                      如果模型在推理的话，tgt = None
# @input_param (pad) : padding 代号，src 和 tgt 中某个元素如果出现 pad 就表明这个位置是 padding
#
# @attr (src) : 就是 输入的 src
# @attr (src_mask) : 根据 src 对应位置是否是 padding 创建出的一个 bool tensor,
#                    其 size : (batch_size, 1, src_len)
# @attr (tgt) : 去掉了最后一列结束符 的 self.tgt
# @attr (tgt_mask) : 根据 tgt 创建的 包含了 Padding Mask 和 Subsequent Mask
#                    结合了使用 padding 来掩盖位置，还防止 train 中模型提前看到答案作弊
#                    它的 size 和 src_mask 一样，都是 (batch_size, 1, src_len)
# @attr (ntokens) :  self.tgt 真正的 token 个数 (不包含 起始符、padding )
#
# @f (make_std_mask) :
# =======================================================================================
class Batch:
    """Object for holding a batch of data with mask during training."""

    # src, tgt : (batch_size, src_len), (batch_size, tgt_len)
    # pad 是因为 一个 batch 中输入的句子长度不一定都一样，所以要进行补齐，比如
    # [1 5 67 32 90] -> [1 5 67 32 90 2 2 2] 其中 1，5，67，32，90都是汉字的代号，2是特殊的填充标记
    # 这个填充标记可能是其他代号，比如 ： 0 可能代表 <pad> 或 <unk>, 1 代表 <start>, 2 代表 <blank>
    # 如果是 推理 的话 ， tgt 可能为 None
    def __init__(self, src, tgt=None, pad=2):  # 2 = <blank>
        self.src = src
        # 所以用 src 和 pad 进行对比就知道 哪些是 padding, 哪些是正常的 文本
        # 然后 在 -2 维度添加一个维度 [batch_size, src_len] -> [batch_size, 1, src_len]
        self.src_mask = (src != pad).unsqueeze(-2)
        if tgt is not None:
            # 比如 tgt 是一个 (batch, tgt_len) ternsor : [[1, 4, 5, 6, 2]]  # 1=<start>, 2=<end>
            # 去掉最后一个 标志 ，通常是结束标志 <end> : [[1, 4, 5, 6]]
            self.tgt = tgt[:, :-1]
            # 去掉第一个 标志，通常是开始标志 <start> : [[4, 5, 6, 2]]
            self.tgt_y = tgt[:, 1:]
            #
            self.tgt_mask = self.make_std_mask(self.tgt, pad)
            # tgt 中 不是 pad符号 的 token 个数
            # TODO: 好像包含了结束符？
            self.ntokens = (self.tgt_y != pad).data.sum()

    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        # 创建一个 padding_mask 用于屏蔽掉 paddings
        # tgt = [[1, 3, 2, 2]] -> [[True, True, False, False]]
        # 在 -2 添加维度是为了和后面 的 feature mask 广播来进行 & 操作
        tgt_mask = (tgt != pad).unsqueeze(-2)
        # 创建一个 feature_mask 用于防止 decoder 看到答案
        tgt_mask = tgt_mask & subsequentMask(*tgt.size()).type_as(
            tgt_mask.data
        )
        return tgt_mask


if __name__ == "__main__":
    pad_id = 2  # <blank>，和 Batch 里默认 pad 一致

    # 模拟一个 batch（2 条句子，每句长度为 6）
    src = torch.tensor([
        [0, 1, 2, 3, 4, pad_id],  # 第一句有一个 padding
        [1, 2, 3, pad_id, pad_id, pad_id],  # 第二句有三个 padding
    ])

    # tgt 中：
    # 假设 0 是 <start>，1-9 是实际 token，10 是 <end>
    tgt = torch.tensor([
        [0, 5, 6, 7, 10, pad_id],  # 有一个 padding
        [0, 3, 4, 10, pad_id, pad_id]  # 有两个 padding
    ])

    # 创建 batch 实例
    batch = Batch(src, tgt, pad=pad_id)

    print("SRC:")
    print(batch.src)

    print("\nSRC_MASK:")
    print(batch.src_mask)  # shape: (batch_size, 1, src_len)

    print("\nTGT:")
    print(batch.tgt)  # 去掉 <end>

    print("\nTGT_Y:")
    print(batch.tgt_y)  # 去掉 <start>

    print("\nTGT_MASK:")
    print(batch.tgt_mask)  # 应该既包含 padding mask 又包含 look-ahead mask

    print("\nNTOKENS:")
    print(batch.ntokens)  # tgt_y 中非 pad 的 token 数量
